import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# 准备时间序列数据
data = np.sin(np.arange(0, 100, 0.1))
seq_length = 20
x = []
y = []
for i in range(len(data) - seq_length - 1):
    x.append(data[i:i + seq_length])
    y.append(data[i + seq_length])
x = np.array(x)
y = np.array(y)
# 将输入数据变为8维
x = np.reshape(x, (x.shape[0], x.shape[1], 1))
x = np.repeat(x, 8, axis=-1)
y = np.reshape(y, (y.shape[0], 1))


# 定义模型
class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(LSTM, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.linear = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        output, _ = self.lstm(x)
        output = self.linear(output[:, -1, :])
        return output


input_size = 8  # 将输入数据变为8维
hidden_size = 32
output_size = 1
learning_rate = 0.01
epochs = 100

model = LSTM(input_size, hidden_size, output_size)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(epochs):
    inputs = torch.from_numpy(x).float()
    targets = torch.from_numpy(y).float()
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    if epoch % 10 == 0:
        print('Epoch: {}, Loss: {:.5f}'.format(epoch, loss.item()))

# 保存模型
torch.save(model.state_dict(), 'lstm_model.pth')

# 加载模型
model = LSTM(input_size, hidden_size, output_size)
model.load_state_dict(torch.load('lstm_model.pth'))

# 预测未来的值
future = 100
x_test = data[-seq_length:]
x_test = np.repeat(x_test, 8)  # 将测试数据变为8维
predictions = []
model.eval()
with torch.no_grad():
    for i in range(future):
        x_input = np.array(x_test[-seq_length:]).reshape((1, seq_length, 8))
        x_input = torch.from_numpy(x_input).float()
        y_output = model(x_input)
        prediction = y_output.detach().numpy().reshape(-1)
        predictions.append(prediction)
        x_test = np.concatenate((x_test, prediction))

# 绘制预测结果
plt.figure(figsize=(12, 6))
plt.plot(np.arange(len(data)), data, label='Original Data')
plt.plot(np.arange(len(data), len(data) + future), predictions, label='Predictions')
plt.legend()
plt.show()
